from typing import Callable

import torch


class MultiLayerPerceptron(torch.nn.Module):
    def __init__(self,
        i_size: int | None,
        *l_sizes: int,
        activation: Callable[[torch.Tensor], torch.Tensor] | None = torch.nn.functional.relu,
        last_activation: Callable[[torch.Tensor], torch.Tensor] | None | bool = None,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ):
        super().__init__()
        if last_activation is True:
            last_activation = activation
        if last_activation is False:
            last_activation = None

        layers = []
        if len(l_sizes) > 0:
            # layers[0]
            o_size = l_sizes[0]
            layer = torch.nn.Linear(i_size, o_size, dtype=dtype, device=device) if i_size is not None \
                else torch.nn.LazyLinear(o_size, dtype=dtype, device=device)
            layers.append(layer)

            # layers[1:]
            for i_size, o_size in zip(l_sizes[:-1], l_sizes[1:]):
                layer = torch.nn.Linear(i_size, o_size, dtype=dtype, device=device)
                layers.append(layer)

        self.activation = activation
        self.last_activation = last_activation
        self.layers = torch.nn.ModuleList(layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # layers[:-1]
        for layer in self.layers[:-1]:
            x = layer(x)
            if self.activation:
                x = self.activation(x)

        # layers[-1]
        x = self.layers[-1](x)
        if self.last_activation:
            x = self.last_activation(x)

        return x
